import os
import sys
import fire
import time

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2,4"

project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
print(project_dir_path)
sys.path.append(project_dir_path)

from omics_bert_label_encoding import process_encoding_embeddings_charls

def main(
    codebook_path: str = os.path.join(project_dir_path, "datasets/ready/charls/v1/codebook_2000_2018.csv"),
    output_dir: str = os.path.join(project_dir_path, "embeddings_local/charls/v1/encoding_embeddings/clhls_2000_2018/"),
    model_name: str = "Qwen/Qwen3-Embedding-8B",
    tensor_parallel_size: int = 2,
    batch_size: int = 8192,
):
  
    start_time = time.time()
    print(f"Start time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}")
    
    process_encoding_embeddings_charls(
        codebook_path=codebook_path,
        output_dir=output_dir,
        model_name=model_name,
        tensor_parallel_size=tensor_parallel_size,
        batch_size=batch_size,
    )

    end_time = time.time()
    print(f"End time: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))}")
    print(f"Total runtime: {(end_time - start_time) / 60} minutes")

if __name__ == "__main__":
    fire.Fire(main)